@coderabbitai summary

Prerequisite

• [None][perf] Optimize MiniMax-M3 CTX block selection #16818

Description

MiniMax-M3 currently projects main Q/K/V and index-Q/index-K separately, then
runs separate normalization, RoPE, quantization, and paged-cache insertion
steps. For long-context prefill, those extra projections and producer launches
are visible in both the critical path and the Nsight Systems trace.

This PR adds an opt-in MiniMax-M3 prefill path that:

• packs Q, K, V, index-Q, and index-K into one MXFP8 projection with the
  MiniMax-M3 TP sharding and replication geometry;
• consumes the packed BF16 output in one horizontal CUDA producer that applies
  Q/K/index-Q/index-K normalization and RoPE, converts the outputs to FP8, and
  directly inserts main K/V plus index-K into their paged caches; and
• returns compact main Q and index-Q tensors to the attention and selector
  paths, avoiding intermediate producer tensors and cache-write launches.

The prototype is disabled by default through
fuse_qkv_index_projection=False. It currently targets pure disaggregated
prefill workers. Decode and mixed workers should leave the option disabled
while their packed-projection behavior is evaluated separately; unsupported
prefill geometry, dtype, and layout combinations retain the compatibility
path.

The completed matched GB300 NVL72 2P8D prefill-isolation A/B used TP2/EP2
attention-DP on the 2-GPU CTX worker, TP8/EP8 attention-DP on the 8-GPU GEN
worker, concurrency 256, exact 8K input, and configured 1-token output. The
baseline was Slurm job 2490293; the candidate was 2490294, with the packed
option enabled only on CTX and disabled on GEN in both runs. Both completed
2,560/2,560 requests:

The client JSON recorded two output tokens per request despite the configured
random_osl=1; this was identical for both variants and does not affect the
relative A/B comparison. The result isolates the prefill benefit; decode
remains explicitly out of scope for this opt-in path.

Test Coverage

• Changed-file pre-commit: passed on all 18 files in the stacked diff.
• Lyris GB300 full and incremental builds: Slurm 2487662, 2487962, and
  2488228 passed.
• Expanded CUDA/Python test job 2488282 passed:
  • horizontal producer numerical tests: 3/3;
  • MiniMax-M3 model tests: 10/10;
  • MSA/index tests: 4/4; and
  • five-way loader geometry and selective-skip tests: 2/2.
• Native checkpoint packing matched exactly for TP2 ranks 0-1 and TP8 ranks
  0-7, including weights, scales, sharding, and replication.
• Real TP4/EP4 attention-DP checkpoint loading completed on all four ranks
  (2090/2090 modules per rank).
• Packaged SQSH smoke job 2488338 passed.
• Full TP4/EP4 MMLU (4096 samples) and GSM8K (1319 samples) registered accuracy
  gates passed in Slurm job 2488379 with the fused projection enabled.
• Matched GB300 NVL72 2P8D prefill-isolation jobs 2490293 and 2490294 completed
  2,560/2,560 requests with exact 8K input and concurrency 256. Enabling the
  packed projection only on CTX improved total token throughput by 8.36% and
  mean TTFT by 7.69%.

PR Checklist

Please review the following before submitting your PR:

• PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

• PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

• Test cases are provided for new code paths (see test instructions)

• If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.

• Any new dependencies have been scanned for license and vulnerabilities

• CODEOWNERS updated if ownership changes

• Documentation updated as needed

• Update tava architecture diagram if there is a significant design change in PR.

• The reviewers assigned automatically/manually are appropriate for the PR.

• Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.